#include "client.h"

// 定义固定的测试数据，便于验证
const std::string EchoClient::TEST_DATA = "test_echo_data";

EchoClient::EchoClient(const char* ip, int port, int numConnections) 
    : ip_(ip), port_(port), numConnections_(numConnections), epollFd_(-1) {}

EchoClient::~EchoClient() {
    if (epollFd_ != -1) close(epollFd_);
}

// 设置socket为非阻塞模式
void EchoClient::setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) {
        std::cerr << "fcntl F_GETFL failed: " << strerror(errno) << std::endl;
        return;
    }
    if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) {
        std::cerr << "fcntl F_SETFL O_NONBLOCK failed: " << strerror(errno) << std::endl;
    }
}

// 检查非阻塞connect是否成功
bool EchoClient::checkConnectSuccess(int fd) {
    int error = 0;
    socklen_t len = sizeof(error);
    if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len) == -1) {
        return false;
    }
    return error == 0;
}

// 完整接收数据
int EchoClient::recvAll(int fd, char* buffer, int expectedLen) {
    int total = 0;
    int retryCount = 0;
    
    // 最多重试5次，避免无限等待
    while (total < expectedLen && retryCount < 5) {
        int n = read(fd, buffer + total, expectedLen - total);
        if (n == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // 暂时没有更多数据，等待10ms后重试
                usleep(1000);
                retryCount++;
                continue;
            }
            return -1; // 错误
        }
        if (n == 0) {
            return total; // 连接关闭
        }
        total += n;
        retryCount = 0; // 成功读取数据，重置重试计数
    }
    return total;
}

// 初始化epoll实例
bool EchoClient::init() {
    epollFd_ = epoll_create1(0);
    if (epollFd_ == -1) {
        std::cerr << "Epoll create failed: " << strerror(errno) << std::endl;
        return false;
    }
    return true;
}

// 创建单个连接并加入epoll
int EchoClient::connectToServer() {
    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd == -1) {
        std::cerr << "Create socket failed: " << strerror(errno) << std::endl;
        return -1;
    }

    setNonBlocking(sockfd);

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port_);
    if (inet_pton(AF_INET, ip_.c_str(), &addr.sin_addr) <= 0) {
        std::cerr << "Inet_pton failed: " << strerror(errno) << std::endl;
        close(sockfd);
        return -1;
    }

    // 非阻塞连接
    int ret = connect(sockfd, (struct sockaddr*)&addr, sizeof(addr));
    if (ret == -1 && errno != EINPROGRESS) {
        std::cerr << "Connect failed: " << strerror(errno) << std::endl;
        close(sockfd);
        return -1;
    }

    // 将socket加入epoll
    struct epoll_event ev;
    ev.events = EPOLLOUT | EPOLLIN | EPOLLERR;
    ev.data.fd = sockfd;
    if (epoll_ctl(epollFd_, EPOLL_CTL_ADD, sockfd, &ev) == -1) {
        std::cerr << "Epoll_ctl add sockfd failed: " << strerror(errno) << std::endl;
        close(sockfd);
        return -1;
    }

    return sockfd;
}

// 批量创建连接
void EchoClient::createConnections() {
    for (int i = 0; i < numConnections_; ++i) {
        int fd = connectToServer();
        if (fd != -1) {
            std::cout << "Created connection " << (i+1) << "/" << numConnections_ << std::endl;
        } else {
            std::cerr << "Failed to create connection " << (i+1) << "/" << numConnections_ << std::endl;
        }
        usleep(100); // 控制连接创建速度
    }
}

// 事件循环
void EchoClient::eventLoop() {
    struct epoll_event events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];
    const char* sendData = "hello";
    int sendLen = strlen(sendData);
    std::atomic<int> failedConnections_{0};

    // 用于缓存部分接收的数据
    std::unordered_map<int, std::pair<std::string, int>> recvBuffer;

    while (true) {
        int nfds = epoll_wait(epollFd_, events, MAX_EVENTS, 1000);
        if (nfds == -1) {
            if (errno == EINTR) continue;
            std::cerr << "Epoll wait failed: " << strerror(errno) << std::endl;
            break;
        }

        for (int i = 0; i < nfds; ++i) {
            int fd = events[i].data.fd;
            uint32_t event = events[i].events;

            // 处理错误事件
            if (event & EPOLLERR) {
                std::cerr << "Connection error on fd: " << fd << std::endl;
                epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                close(fd);
                failedConnections_++;
                recvBuffer.erase(fd);
                continue;
            }

            // 处理连接完成（写事件）
            if (event & EPOLLOUT) {
                // 检查连接是否成功
                if (!checkConnectSuccess(fd)) {
                    std::cerr << "Connection failed on fd: " << fd << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    continue;
                }

                // 连接成功，发送测试数据
                connected_++;
                std::cout << "Connected (" << connected_ << "/" << numConnections_ << ") on fd: " << fd << std::endl;

                // 发送固定测试数据
                int nwrite = write(fd, sendData, sendLen);
                if (nwrite == -1) {
                    if (errno == EAGAIN || errno == EWOULDBLOCK) {
                        // 暂时无法发送，保留事件
                        continue;
                    }
                    std::cerr << "Write failed on fd " << fd << ": " << strerror(errno) << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    continue;
                }

                // std::cout << "Client sent to fd " << fd << ": '" << TEST_DATA << "' (" << nwrite << " bytes)" << std::endl;

                // 改为监听读事件，等待回显
                struct epoll_event ev;
                ev.events = EPOLLIN | EPOLLERR;
                ev.data.fd = fd;
                if (epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev) == -1) {
                    std::cerr << "Epoll modify failed on fd " << fd << std::endl;
                    close(fd);
                    failedConnections_++;
                }
            }

            // 处理服务器响应（读事件）
            if (event & EPOLLIN) {
                // 初始化接收缓存
                if (recvBuffer.find(fd) == recvBuffer.end()) {
                    recvBuffer[fd] = std::make_pair("", 0);
                }
                std::string& data = recvBuffer[fd].first;
                int& received = recvBuffer[fd].second;

                // 接收数据直到完整
                int nbytes = recvAll(fd, buffer, sendLen - received);
                if (nbytes == -1) {
                    std::cerr << "Read failed on fd " << fd << ": " << strerror(errno) << std::endl;
                    recvBuffer.erase(fd);
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    continue;
                }

                // 更新接收缓存
                received += nbytes;
                data.append(buffer, nbytes);
                std::cout << "Client received from fd " << fd << ": '" << data << "' (" << received << "/" << sendLen << " bytes)" << std::endl;

                // 检查是否接收完整
                if (received == sendLen) {
                    // 验证回显数据
                    std::string cmp(sendData);
                    if (data == cmp) {
                        echoed_++;
                        std::cout << "Echo success (" << echoed_ << "/" << numConnections_ << ") on fd: " << fd << std::endl;
                    } else {
                        std::cerr << "Echo mismatch on fd " << fd << ": expected '" << sendData 
                                  << "', got '" << data << "'" << std::endl;
                        failedConnections_++;
                    }

                    // 清理资源
                    recvBuffer.erase(fd);
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                } else if (received > sendLen) {
                    // 数据长度异常
                    std::cerr << "Unexpected data length on fd " << fd << ": got " << received << ", expected " << sendLen << std::endl;
                    recvBuffer.erase(fd);
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                }
                // 否则继续等待更多数据
            }
        }

        // 所有连接处理完成则退出
        if (echoed_ + failedConnections_ >= numConnections_) {
            break;
        }
    }

    // 输出最终统计
    std::cout << "===== Test Result =====" << std::endl;
    std::cout << "Total connections: " << numConnections_ << std::endl;
    std::cout << "Successfully connected: " << connected_ << std::endl;
    std::cout << "Successfully echoed: " << echoed_ << std::endl;
    std::cout << "Failed connections: " << failedConnections_ << std::endl;
}

// 启动客户端
void EchoClient::start() {
    std::thread connThread(&EchoClient::createConnections, this);
    eventLoop();
    connThread.join();
}
